// hdr.cpp
//

#include "loadPNM.h"
#include "PFMImage.h"

#include <vector>
#include <math.h>
#include <string>

#define DEBUG

#ifdef DEBUG 
#define DPRINTF(fmt, ...) \
    do { printf(__FILE__": " fmt, ## __VA_ARGS__); } while (0)
#else
#define DPRINTF(fmt, ...) \
    do { } while (0)
#endif

using namespace std;

vector <string> PFM_FILES = {
		"Memorial/memorial1.pfm",
        "Memorial/memorial2.pfm",
        "Memorial/memorial3.pfm",
        "Memorial/memorial4.pfm",
        "Memorial/memorial5.pfm",
        "Memorial/memorial6.pfm",
        "Memorial/memorial7.pfm" };

vector <unsigned> EXPOSURE = { 1, 4, 16, 64, 256, 1024, 4096 };

void loadPFMs (vector <PFMImage*> & pfms, vector <string> & files)
{
    for (string& filename : files)
    {
        PFMImage* ppfm = new PFMImage(filename);
        pfms.push_back(ppfm); // Remember to delete these pointers
    }
    return;
}

float centerWeighted (float x)
{
    // Simple triangle function
    return fmax(1 - fabs(x-0.5)*2, 0);
}

void combineImages (vector <PFMImage*> & pfms, PFMImage & result)
{
    // Assume all images share the same characteristics
	PFMImage* front = pfms.front();
    uint width = front->getWidth();
    uint height = front->getHeight();
    uint num_components = front->getNumComponents();

    uint num_pixels = width * height;

    // Retrieve a pointer to each images data buffer & the rel. exposure
    float* imgDataPtrs[pfms.size()];
    unsigned exposures[pfms.size()];

    for (int i=0; i<pfms.size(); i++){
        imgDataPtrs[i] = pfms.at(i)->getData();
        exposures[i] = pfms.at(i)->getExposure();
    }

    // Retrieve a pointer to the target image's data
    float* targetData = result.getData();

	// For each pixel
	for (int p = 0; p < num_pixels; p++)
	{
		uint pixelOffset = p * num_components;

		// For each channel
		for (int c = 0; c < num_components; c++)
		{
			float sum1 = 0, sum2 = 0;

			// For each input PFM image
			for (int i = 0; i < pfms.size(); i++)
			{

				// Limit the Zi values to this range
				// (skip the whole image's pixel to avoid color shifting)
				if (c == 0)
				{
					float avg_zi = 0;

					for (int k = 0; k < num_components; k++)
					{
						avg_zi += imgDataPtrs[i][pixelOffset + k];
					}

					avg_zi /= num_components;

					if (avg_zi < 0.005 || avg_zi > 0.92)
						// Goto next image
						continue;
				}

				// Get the input value
				float zi = imgDataPtrs[i][pixelOffset + c];

				//DPRINTF("IMG:%d Pixel pos (%d, %d):%d, Zi=%f\n", i, x, y, c, zi);

				float ex = (float) exposures[i];
				sum1 += log((1.0 / ex) * zi) * centerWeighted(zi);
				sum2 += centerWeighted(zi);
			}

			/* F(x,y) = exp(sum_i(log((1/Ei)*Zi(x,y))*w(Zi(x,y))) /
			 *  sum_i(w(Zi(x,y)))) */

			targetData[pixelOffset + c] = exp(sum1 / sum2);
		}
	}
}


float min_radiance(PFMImage& img)
{
    float* imgData = img.getData();
    uint width = img.getWidth();
    uint height = img.getHeight();
    uint num_components = img.getNumComponents();

	uint num_pixels = width * height;


    float min_avg = 1.0;

    for (int p=0; p<num_pixels; p++ ) {
    	uint pixelOffset = p * num_components;

    	float avg = 0;

    	for (int c = 0 ; c < num_components ; c++) {
    		avg += imgData[pixelOffset + c];
    	}

    	avg /= num_components;
    	min_avg = fmin(min_avg, avg);
    }

    return min_avg;
}

float max_radiance(PFMImage& img)
{
    float* imgData = img.getData();
    uint width = img.getWidth();
    uint height = img.getHeight();
    uint num_components = img.getNumComponents();

	uint num_pixels = width * height;


    float max_avg = 0;

    for (int p=0; p<num_pixels; p++ ) {
    	uint pixelOffset = p * num_components;

    	float avg = 0;

    	for (int c = 0 ; c < num_components ; c++) {
    		avg += imgData[pixelOffset + c];
    	}

    	avg /= num_components;
    	max_avg = fmax(max_avg, avg);
    }

    return max_avg;
}


void tone_map(PFMImage& img)
{
	float* imgData = img.getData();
	uint width = img.getWidth();
	uint height = img.getHeight();
	uint num_components = img.getNumComponents();

	uint num_pixels = width * height;

	// Get the maximum radiance value
	float max_f = max_radiance(img);

	for (int p=0; p<num_pixels; p++ )
	{
		uint pixelOffset = p * num_components;

		for (int c = 0; c < num_components; c++)
		{
			imgData[pixelOffset + c] /= max_f;
		}
	}
}

void increase_exposure(PFMImage& img, uint num_stops=1)
{
	float* imgData = img.getData();
	uint width = img.getWidth();
	uint height = img.getHeight();
	uint num_components = img.getNumComponents();

	uint num_pixels = width * height;

	for (int p=0; p<num_pixels; p++ ) {

		uint pixelOffset = p * num_components;

		for (int c = 0; c < num_components; c++)
		{
			// Increase by 1 stop
			imgData[pixelOffset + c] *= pow(2, num_stops);

			// Clamp to 1.0
			float z = imgData[pixelOffset + c];
			imgData[pixelOffset + c] = z > 1.0 ? 1.0 : z;
		}
	}
}

void gamma(PFMImage& img, float gamma=1.5)
{
	float* imgData = img.getData();
	uint width = img.getWidth();
	uint height = img.getHeight();
	uint num_components = img.getNumComponents();

	uint num_pixels = width * height;

	for (int p=0; p<num_pixels; p++ ) {

		uint pixelOffset = p * num_components;

		for (int c = 0; c < num_components; c++)
		{
			// Increase by 1 stop
			float z = imgData[pixelOffset + c];
			imgData[pixelOffset + c] = pow(z, 1/gamma);
		}
	}
}


int main (int argc, char * * argv)
{
    vector<PFMImage*> pfms;
    loadPFMs(pfms, PFM_FILES);

    // Assign the relative exposure to each loaded image
    for ( int i = 0; i < pfms.size(); i++ )
    {
        int exp = EXPOSURE.at(i);
        PFMImage* pfm = pfms.at(i);
        pfm->setExposure(exp);
    }

    bool coursework_mode = false;

    // Default values as found by trial and error
    uint num_stops = 8;
    float gamma_val = 2.2;
    char* output_filename = "g";

    switch(argc)
    {
    case 4:
    	output_filename = argv[3];
    	/* no break */
    case 3:
    	gamma_val = atof(argv[2]);
    	/* no break */
    case 2:
    	num_stops = atoi(argv[1]);
    	break;

    default:
    	coursework_mode = true;
    	break;
    };

    // Create an empty PFM image with the same size as one of the given PFMs
    PFMImage* first = pfms.front();
    unsigned w = first->getWidth();
    unsigned h = first->getHeight();
    unsigned nc = first->getNumComponents();

    PFMImage hdr_image = PFMImage(w, h, nc);

    // Combine the images into a single HDR PFM image
    combineImages(pfms, hdr_image);

	// Apply post-processing operations
    if (!coursework_mode) {
    	tone_map(hdr_image);
    	increase_exposure(hdr_image, num_stops);
    	gamma(hdr_image, gamma_val);
    }
    else {
		hdr_image.writePPM("f");
    	tone_map(hdr_image);
		hdr_image.writePPM("n");
    	increase_exposure(hdr_image, num_stops);
    	gamma(hdr_image, gamma_val);

    	float drange = max_radiance(hdr_image) / min_radiance(hdr_image);
    	DPRINTF("Dynamic range of the image: %f\n", drange);
    }

    // Save the final image
    hdr_image.writePPM(output_filename);

    // Cleanup
    for ( auto& it : pfms )
        delete it;

    return 0;
}
